package com.buffer.springai.service;

import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.Role;
import com.buffer.springai.config.NacosConfig;
import io.reactivex.Flowable;
import io.reactivex.schedulers.Schedulers;
import org.springframework.stereotype.Service;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import javax.annotation.Resource;
import java.util.Arrays;

@Service
public class StreamService {

    @Resource
    private NacosConfig nacosConfig;

    public void askAIWithStream(String question, SseEmitter emitter) {
        // 初始化 Generation 实例
        Generation gen = new Generation();

        // 构建请求参数
        GenerationParam param = GenerationParam.builder()
                .apiKey(nacosConfig.getApiKey())
                .model("qwen-flash")
                .messages(Arrays.asList(
                        Message.builder()
                                .role(Role.USER.getValue())
                                .content(question)
                                .build()
                ))
                .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                .incrementalOutput(true) // 开启增量输出，流式返回
                .build();
                
        // 发起流式调用并处理响应
        try {
            Flowable<GenerationResult> result = gen.streamCall(param);
            // 用于跟踪之前发送的内容
            StringBuilder sentContent = new StringBuilder();
            
            result
                    .subscribeOn(Schedulers.io()) // IO线程执行请求
                    .observeOn(Schedulers.computation()) // 计算线程处理响应
                    .subscribe(
                            // onNext: 处理每个响应片段
                            message -> {
                                try {
                                    String fullContent = message.getOutput().getChoices().get(0).getMessage().getContent();
                                    String finishReason = message.getOutput().getChoices().get(0).getFinishReason();
                                    emitter.send(fullContent);
                                    
//                                    // 检查是否有新内容需要发送
//                                    if (fullContent != null && fullContent.length() > sentContent.length()) {
//                                        // 获取新增的内容
//                                        String newContent = fullContent.substring(sentContent.length());
//                                        // 更新已发送内容的跟踪
//                                        sentContent.append(newContent);
//                                        // 发送新增内容到前端
//                                        emitter.send(SseEmitter.event().data(newContent));
//                                    }
                                    
                                    // 当 finishReason 不为 null 时，表示是最后一个 chunk，输出用量信息
                                    if (finishReason != null && !"null".equals(finishReason)) {
                                        String usageInfo = "\n--- 请求用量 ---\n" +
                                                "输入 Tokens：" + message.getUsage().getInputTokens() + "\n" +
                                                "输出 Tokens：" + message.getUsage().getOutputTokens() + "\n" +
                                                "总 Tokens：" + message.getUsage().getTotalTokens();
                                        emitter.send(SseEmitter.event().data(usageInfo));
                                    }
                                } catch (Exception e) {
                                    emitter.completeWithError(e);
                                }
                            },
                            // onError: 处理错误
                            error -> {
                                try {
                                    emitter.send(SseEmitter.event().data("Error: " + error.getMessage()));
                                } catch (Exception e) {
                                    // 忽略发送错误
                                }
                                emitter.completeWithError(error);
                            },
                            // onComplete: 完成回调
                            () -> {
                                emitter.complete();
                            }
                    );
        } catch (Exception e) {
            try {
                emitter.completeWithError(e);
            } catch (Exception ex) {
                // 忽略内部异常
            }
        }
    }
}